import time
import pypdfium2
import os
import argparse
from marker.convert import convert_single_pdf
from marker.logger import configure_logging
from marker.models import load_all_models
from marker.output import save_markdown
from joblib import dump, load
from concurrent.futures import ThreadPoolExecutor, as_completed

class PDFConverter:
    # ... (保持其他方法不变)

    def _process_single_file(self, input_file, output_folder):
        if self.model_lst is None:
            print('加载模型')
            self.model_lst = self._load_models()
        start = time.time()
        try:
            # 尝试打开PDF文件
            pdf = pypdfium2.PdfDocument(input_file)

            full_text, images, out_meta = convert_single_pdf(
                input_file,
                self.model_lst,
                max_pages=self.max_pages,
                langs=self.langs,
                batch_multiplier=self.batch_multiplier,
                start_page=self.start_page,
                ocr_all_pages=self.ocr_all_pages
            )

            fname = os.path.basename(input_file)
            subfolder_path = save_markdown(output_folder, fname, full_text, images, out_meta)

            print(f"Saved markdown to the {subfolder_path} folder")
            print(f"Processing time: {time.time() - start:.2f} seconds")

            return subfolder_path
        except Exception as e:
            print(f"Error processing file {input_file}: {str(e)}")
            return None

    def process_directory(self, input_dir, output_dir, max_workers=4):
        if not os.path.exists(input_dir):
            raise FileNotFoundError(f"Input directory not found: {input_dir}")

        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        pdf_files = []
        for root, _, files in os.walk(input_dir):
            for file in files:
                if file.lower().endswith('.pdf'):
                    pdf_files.append(os.path.join(root, file))

        if not pdf_files:
            print(f"No PDF files found in {input_dir}")
            return

        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = []
            for pdf_file in pdf_files:
                relative_path = os.path.relpath(pdf_file, input_dir)
                output_subdir = os.path.join(output_dir, os.path.dirname(relative_path))
                print('正在处理', pdf_file)
                futures.append(executor.submit(self.process_file, pdf_file, output_subdir))

            for future in as_completed(futures):
                try:
                    subfolder_path = future.result()
                    if subfolder_path:
                        print(f"Processed: {subfolder_path}")
                    else:
                        print("Failed to process file")
                except Exception as e:
                    print(f"An error occurred: {str(e)}")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("input", help="PDF file or directory to parse")
    parser.add_argument("output", help="Output base folder path")
    parser.add_argument("--is_directory_input", type=int, default=0, help="是否为文件夹")
    parser.add_argument("--max_pages", type=int, default=None, help="Maximum number of pages to parse")
    parser.add_argument("--start_page", type=int, default=None, help="Page to start processing at")
    parser.add_argument("--langs", type=str, help="Optional languages to use for OCR, comma separated", default=None)
    parser.add_argument("--batch_multiplier", type=int, default=2, help="How much to increase batch sizes")
    parser.add_argument("--debug", action="store_true", help="Enable debug logging", default=False)
    parser.add_argument("--ocr_all_pages", action="store_true", help="Force OCR on all pages", default=False)
    parser.add_argument("--max_workers", type=int, default=1, help="Maximum number of worker threads for batch processing")
    args = parser.parse_args()

    configure_logging()

    langs = args.langs.split(",") if args.langs else None

    converter = PDFConverter(
        max_pages=args.max_pages,
        langs=langs,
        batch_multiplier=args.batch_multiplier,
        start_page=args.start_page,
        ocr_all_pages=args.ocr_all_pages
    )

    try:
        if args.is_directory_input or os.path.isdir(args.input):
            converter.process_directory(args.input, args.output, max_workers=args.max_workers)
        else:
            subfolder_path = converter.process_file(args.input, args.output)
            if subfolder_path:
                print(f"Processed: {subfolder_path}")
            else:
                print("Failed to process file")
        if args.debug:
            print(f"Total time: {time.time() - start}")
    except Exception as e:
        print(f"An error occurred: {str(e)}")

if __name__ == "__main__":
    main()
